#!/usr/bin/env python3
"""
Phase 4A — 45s Pilot 合成脚本
1. 生成 TTS 口播 (edge-tts)
2. 生成 SRT 字幕
3. 生成 BGM + SFX
4. 拼接 5 scene clips + 混音 + 字幕 = 最终 Pilot
"""

import os, sys, subprocess, json, math

ROOT = r"E:\集群文件夹\factory_os\short_video_real_data_pipeline\phase4a_45s_pilot"
CLIPS_DIR    = os.path.join(ROOT, "03_scene_clips")
VO_DIR       = os.path.join(ROOT, "04_voiceover")
SUB_DIR      = os.path.join(ROOT, "05_subtitles")
SFX_DIR      = os.path.join(ROOT, "06_music_sfx")
EXPORT_DIR   = os.path.join(ROOT, "07_exports")

WIDTH, HEIGHT = 1080, 1920
FPS = 15

# ── Scene 定义 (顺序、时长、TTS文本) ──
SCENES = [
    {
        "name": "scene_01_hook_error",
        "clip": "scene_01_hook_error.mp4",
        "duration": 8.0,
        "tts": "两千条数据，差点全错。没有真实链接、没有作者信息、没有互动指标——这样的数据，根本不能用来做对标分析。",
    },
    {
        "name": "scene_02_data_gate",
        "clip": "scene_02_data_gate.mp4",
        "duration": 10.0,
        "tts": "所以我们做了真实数据质量门禁。十条关键字段逐一检查：标题、链接、作者、时间、互动数据……全部通过，才能进入候选池。",
    },
    {
        "name": "scene_03_candidates",
        "clip": "scene_03_candidates.mp4",
        "duration": 10.0,
        "tts": "最终筛选出八十九条真实候选。每一条都有完整的视频链接，每一条都是真实的抖音创作者，每一条都有可验证的互动数据。",
    },
    {
        "name": "scene_04_download_success",
        "clip": "scene_04_download_success.mp4",
        "duration": 10.0,
        "tts": "排名前十的视频本体全部下载成功。我们只做分析，不搬运内容——这是底线。",
    },
    {
        "name": "scene_05_cta",
        "clip": "scene_05_cta.mp4",
        "duration": 7.0,
        "tts": "真实数据先行，不靠假数据骗自己。用真实数据，做有效对标。",
    },
]

TOTAL_DURATION = sum(s["duration"] for s in SCENES)  # ~45s


def ensure_dirs():
    for d in [VO_DIR, SUB_DIR, SFX_DIR, EXPORT_DIR]:
        os.makedirs(d, exist_ok=True)


def step1_generate_tts():
    """用 edge-tts 生成每段口播"""
    print("=" * 60)
    print("Step 1: 生成 TTS 口播")
    print("=" * 60)

    all_ok = True
    for s in SCENES:
        output = os.path.join(VO_DIR, f"{s['name']}.mp3")
        if os.path.exists(output):
            size = os.path.getsize(output) / 1024
            print(f"  [跳过] {s['name']} (已存在 {size:.0f} KB)")
            continue

        print(f"  [生成] {s['name']}...", end=" ", flush=True)
        # Use edge-tts with Chinese female voice
        cmd = [
            sys.executable, "-m", "edge_tts",
            "--voice", "zh-CN-XiaoxiaoNeural",
            "--text", s["tts"],
            "--write-media", output,
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        if result.returncode == 0 and os.path.exists(output):
            size = os.path.getsize(output) / 1024
            print(f"✅ {size:.0f} KB")
        else:
            print(f"❌ {result.stderr}")
            all_ok = False

    return all_ok


def step2_get_tts_durations():
    """获取每段 TTS 的实际时长 (秒)"""
    durations = []
    for s in SCENES:
        mp3 = os.path.join(VO_DIR, f"{s['name']}.mp3")
        if not os.path.exists(mp3):
            durations.append(s["duration"])
            continue
        # Use ffprobe to get duration
        cmd = [
            "ffprobe", "-v", "error",
            "-show_entries", "format=duration",
            "-of", "default=noprint_wrappers=1:nokey=1",
            mp3,
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
        try:
            dur = float(result.stdout.strip())
        except:
            dur = s["duration"]
        durations.append(dur)
        print(f"  {s['name']}: TTS={dur:.1f}s  scene={s['duration']}s")
    return durations


def step3_generate_subtitles(tts_durations):
    """生成 SRT 字幕文件"""
    print("\n" + "=" * 60)
    print("Step 2: 生成字幕")
    print("=" * 60)

    # Use scene durations for subtitle timing
    srt_path = os.path.join(SUB_DIR, "pilot_subtitles.srt")
    lines = []
    idx = 1
    current_time = 0.0

    for i, s in enumerate(SCENES):
        start = current_time
        end = start + s["duration"]

        def fmt(sec):
            h = int(sec // 3600)
            m = int((sec % 3600) // 60)
            s_val = sec % 60
            return f"{h:02d}:{m:02d}:{s_val:06.3f}".replace(".", ",")

        # Split TTS text into subtitle chunks (by punctuation)
        text = s["tts"]
        # For simplicity, use full text as one subtitle per scene
        lines.append(f"{idx}")
        lines.append(f"{fmt(start)} --> {fmt(end)}")
        lines.append(f"{text}")
        lines.append("")
        idx += 1

        current_time = end

    with open(srt_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))
    print(f"  [生成] {srt_path} ({idx-1} 条字幕)")
    return srt_path


def step4_generate_sfx():
    """生成 BGM 和音效"""
    print("\n" + "=" * 60)
    print("Step 3: 生成 BGM + SFX")
    print("=" * 60)

    total_dur = TOTAL_DURATION

    # BGM: 轻柔持续的 pad 音色（低频正弦波 + 谐波）
    bgm_path = os.path.join(SFX_DIR, "bgm.mp3")
    if not os.path.exists(bgm_path):
        print("  [生成] BGM...")
        # Simple ambient: layered sine waves at different frequencies
        # Note: volume is applied via filter, not inside sine source
        cmd = [
            "ffmpeg", "-y",
            "-f", "lavfi", "-i", "sine=frequency=220:duration={}".format(total_dur),
            "-f", "lavfi", "-i", "sine=frequency=330:duration={}".format(total_dur),
            "-f", "lavfi", "-i", "sine=frequency=440:duration={}".format(total_dur),
            "-filter_complex", "[0]volume=0.08[a];[1]volume=0.04[b];[2]volume=0.02[c];[a][b][c]amix=inputs=3:duration=first",
            "-ac", "2",
            "-ar", "44100",
            "-b:a", "128k",
            bgm_path,
        ]
        subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        size = os.path.getsize(bgm_path) / 1024
        print(f"  ✅ BGM: {size:.0f} KB")

    # SFX: 各个音效
    sfx_files = []

    # Error beep (0.3s) - for scene 01
    beep_path = os.path.join(SFX_DIR, "sfx_error_beep.mp3")
    if not os.path.exists(beep_path):
        cmd = [
            "ffmpeg", "-y",
            "-f", "lavfi", "-i",
            "sine=frequency=880:duration=0.15,volume=0.4",
            "-ac", "2", "-ar", "44100", beep_path,
        ]
        subprocess.run(cmd, capture_output=True, text=True, timeout=30)
    sfx_files.append({"path": beep_path, "start": 2.0, "label": "error_beep"})

    # Click (0.1s) - for each check in scene 02
    click_path = os.path.join(SFX_DIR, "sfx_click.mp3")
    if not os.path.exists(click_path):
        cmd = [
            "ffmpeg", "-y",
            "-f", "lavfi", "-i",
            "sine=frequency=1200:duration=0.05,volume=0.3",
            "-ac", "2", "-ar", "44100", click_path,
        ]
        subprocess.run(cmd, capture_output=True, text=True, timeout=30)
    sfx_files.append({"path": click_path, "start": 9.0, "label": "click_1"})
    sfx_files.append({"path": click_path, "start": 9.8, "label": "click_2"})
    sfx_files.append({"path": click_path, "start": 10.6, "label": "click_3"})
    sfx_files.append({"path": click_path, "start": 11.4, "label": "click_4"})
    sfx_files.append({"path": click_path, "start": 12.2, "label": "click_5"})

    # Whoosh (0.5s) - transition
    whoosh_path = os.path.join(SFX_DIR, "sfx_whoosh.mp3")
    if not os.path.exists(whoosh_path):
        cmd = [
            "ffmpeg", "-y",
            "-f", "lavfi", "-i",
            "sine=frequency=200:duration=0.3,volume=0.2",
            "-ac", "2", "-ar", "44100", whoosh_path,
        ]
        subprocess.run(cmd, capture_output=True, text=True, timeout=30)
    sfx_files.append({"path": whoosh_path, "start": 8.0, "label": "whoosh_1"})
    sfx_files.append({"path": whoosh_path, "start": 18.0, "label": "whoosh_2"})
    sfx_files.append({"path": whoosh_path, "start": 28.0, "label": "whoosh_3"})
    sfx_files.append({"path": whoosh_path, "start": 38.0, "label": "whoosh_4"})

    # Success chime (0.5s) - for scene 04 download complete
    chime_path = os.path.join(SFX_DIR, "sfx_chime.mp3")
    if not os.path.exists(chime_path):
        cmd = [
            "ffmpeg", "-y",
            "-f", "lavfi", "-i",
            "sine=frequency=523:duration=0.3,volume=0.35",
            "-ac", "2", "-ar", "44100", chime_path,
        ]
        subprocess.run(cmd, capture_output=True, text=True, timeout=30)
    sfx_files.append({"path": chime_path, "start": 37.0, "label": "chime"})
    sfx_files.append({"path": chime_path, "start": 44.0, "label": "final_chime"})

    print(f"  ✅ {len(sfx_files)} 个音效")

    # Merge all SFX into one track
    merged_sfx = os.path.join(SFX_DIR, "sfx_all.mp3")
    sfx_filter = ""
    sfx_inputs = ""
    total_dur_ms = int(total_dur * 1000)

    # Generate silent base
    silent_path = os.path.join(SFX_DIR, "_silent.mp3")
    if not os.path.exists(silent_path):
        cmd = [
            "ffmpeg", "-y",
            "-f", "lavfi", "-i", "anullsrc=r=44100:cl=mono",
            "-t", str(total_dur),
            "-ac", "2", "-ar", "44100", silent_path,
        ]
        subprocess.run(cmd, capture_output=True, text=True, timeout=30)

    # Use adelay to position each SFX
    # Build complex filter
    filter_parts = []
    input_idx = 0
    input_files = [silent_path]

    # Group by unique file to avoid duplicate inputs
    seen_paths = set()
    for sfx in sfx_files:
        if sfx["path"] not in seen_paths:
            seen_paths.add(sfx["path"])
            input_files.append(sfx["path"])

    # Build filter: [1]adelay=2000[beep]; [2]adelay=9000[click1]; ...
    filter_chain = []
    for i, sfx in enumerate(sfx_files):
        file_idx = input_files.index(sfx["path"])
        delay_ms = int(sfx["start"] * 1000)
        filter_chain.append(f"[{file_idx}]adelay={delay_ms}|{delay_ms}[s{i}]")

    # Mix all
    mix_inputs = "".join(f"[s{i}]" for i in range(len(sfx_files)))
    filter_chain.append(f"{mix_inputs}amix=inputs={len(sfx_files)}:duration=first[sfxout]")

    if not os.path.exists(merged_sfx):
        print("  [合成] 合并所有音效...")

        # Convert timestamps for FFmpeg: force each SFX at its position
        # Simpler approach: use the silent base and overlay each SFX
        # Build a filter that places each sound
        inputs_str = " ".join(f'-i "{p}"' for p in input_files)

        # Simple approach: just use amix with each delayed separately
        filter_graph = "; ".join(filter_chain)

        cmd = [
            "ffmpeg", "-y",
        ] + [item for pair in zip(["-i"] * len(input_files), input_files) for item in pair] + [
            "-filter_complex", filter_graph,
            "-map", "[sfxout]",
            "-ac", "2", "-ar", "44100",
            "-b:a", "128k",
            merged_sfx,
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        if result.returncode == 0:
            size = os.path.getsize(merged_sfx) / 1024
            print(f"  ✅ SFX 合并: {size:.0f} KB")
        else:
            print(f"  ❌ SFX 合并失败: {result.stderr[-200:]}")
            # Fallback: just copy BGM
            import shutil
            shutil.copy(bgm_path, merged_sfx)

    return bgm_path, merged_sfx


def step5_compose_final(bgm_path, sfx_path):
    """合成最终 45s Pilot"""
    print("\n" + "=" * 60)
    print("Step 4: 合成最终 45s Pilot")
    print("=" * 60)

    # Create concat file for video clips
    concat_file = os.path.join(ROOT, "_concat.txt")
    with open(concat_file, "w", encoding="utf-8") as f:
        for s in SCENES:
            clip_path = os.path.join(CLIPS_DIR, s["clip"])
            if os.path.exists(clip_path):
                f.write(f"file '{clip_path}'\n")
                f.write(f"duration {s['duration']}\n")
            else:
                print(f"  [警告] clip 不存在: {clip_path}")

    # TTS concat
    tts_concat = os.path.join(ROOT, "_tts_concat.txt")
    with open(tts_concat, "w", encoding="utf-8") as f:
        for s in SCENES:
            tts_path = os.path.join(VO_DIR, f"{s['name']}.mp3")
            if os.path.exists(tts_path):
                f.write(f"file '{tts_path}'\n")

    # Concat TTS into one file
    tts_all = os.path.join(VO_DIR, "tts_all.mp3")
    if not os.path.exists(tts_all):
        print("  [TTS] 拼接口播...")
        cmd = [
            "ffmpeg", "-y",
            "-f", "concat", "-safe", "0",
            "-i", tts_concat,
            "-ac", "2", "-ar", "44100",
            "-b:a", "128k",
            tts_all,
        ]
        subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        size = os.path.getsize(tts_all) / 1024
        print(f"  ✅ TTS 拼接: {size:.0f} KB")

    # Concat video clips
    video_concat = os.path.join(ROOT, "_video_concat.ts")
    if not os.path.exists(video_concat):
        print("  [视频] 拼接 clips...")
        cmd = [
            "ffmpeg", "-y",
            "-f", "concat", "-safe", "0",
            "-i", concat_file,
            "-c", "copy",
            "-an",
            video_concat,
        ]
        subprocess.run(cmd, capture_output=True, text=True, timeout=120)

    # Determine TTS duration
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        tts_all,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
    try:
        tts_dur = float(result.stdout.strip())
    except:
        tts_dur = TOTAL_DURATION
    print(f"  TTS 时长: {tts_dur:.1f}s  视频时长: {TOTAL_DURATION:.1f}s")

    # Determine bgm duration
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        bgm_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
    try:
        bgm_dur = float(result.stdout.strip())
    except:
        bgm_dur = TOTAL_DURATION

    # If bgm is shorter than video, pad it
    if bgm_dur < TOTAL_DURATION:
        bgm_path_orig = bgm_path
        bgm_path = os.path.join(SFX_DIR, "bgm_padded.mp3")
        cmd = [
            "ffmpeg", "-y",
            "-i", bgm_path_orig,
            "-f", "lavfi", "-t", str(TOTAL_DURATION), "-i", "anullsrc",
            "-filter_complex", "[0:a][1:a]amix=inputs=2:duration=first",
            "-ac", "2", "-ar", "44100",
            bgm_path,
        ]
        subprocess.run(cmd, capture_output=True, text=True, timeout=30)

    # Determine sfx duration
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        sfx_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
    try:
        sfx_dur = float(result.stdout.strip())
    except:
        sfx_dur = TOTAL_DURATION

    # Final mix: video + TTS + BGM + SFX + subtitles
    output = os.path.join(EXPORT_DIR, "pilot_45s_v1.mp4")

    print("  [最终合成] 视频 + TTS + BGM + SFX + 字幕...")

    # Generate ASS subtitle file in C:\Users\Admin (simple path, no special chars for FFmpeg)
    ass_path = r"C:\Users\Admin\pilot_sub.ass"
    with open(ass_path, "w", encoding="utf-8") as f:
        f.write("""[Script Info]
; Generated for Phase 4A Pilot
ScriptType: v4.00+
PlayResX: 1080
PlayResY: 1920
ScaledBorderAndShadow: yes

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Default,Microsoft YaHei,36,&H00E0E0E0,&H000000FF,&H00000000,&H80000000,0,0,0,0,100,100,0,0,1,2,0,2,60,60,140,1

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
Dialogue: 0,0:00:00.00,0:00:08.00,Default,,0,0,0,,两千条数据，差点全错。没有真实链接、没有作者信息、没有互动指标——这样的数据，根本不能用来做对标分析。
Dialogue: 0,0:00:08.00,0:00:18.00,Default,,0,0,0,,所以我们做了真实数据质量门禁。十条关键字段逐一检查：标题、链接、作者、时间、互动数据……全部通过，才能进入候选池。
Dialogue: 0,0:00:18.00,0:00:28.00,Default,,0,0,0,,最终筛选出八十九条真实候选。每一条都有完整的视频链接，每一条都是真实的抖音创作者，每一条都有可验证的互动数据。
Dialogue: 0,0:00:28.00,0:00:38.00,Default,,0,0,0,,排名前十的视频本体全部下载成功。我们只做分析，不搬运内容——这是底线。
Dialogue: 0,0:00:38.00,0:00:45.00,Default,,0,0,0,,真实数据先行，不靠假数据骗自己。用真实数据，做有效对标。
""")

    # Use amix for audio mixing, ass filter for subtitles (relative path avoids colon issues)
    # Run FFmpeg from C:\Users\Admin so ass=pilot_sub.ass resolves as relative path
    cmd = [
        "ffmpeg", "-y",
        "-i", video_concat,
        "-i", tts_all,
        "-i", bgm_path,
        "-i", sfx_path,
        "-filter_complex",
        "[1:a]volume=1.5[a_tts];"
        "[2:a]volume=0.3[a_bgm];"
        "[3:a]volume=0.5[a_sfx];"
        "[a_tts][a_bgm][a_sfx]amix=inputs=3:duration=first:weights=1.5 0.3 0.5[aout]",
        "-map", "0:v",
        "-map", "[aout]",
        "-c:v", "libx264",
        "-preset", "medium",
        "-crf", "18",
        "-c:a", "aac",
        "-b:a", "192k",
        "-vf", "ass=pilot_sub.ass",
        "-pix_fmt", "yuv420p",
        "-movflags", "+faststart",
        output,
    ]

    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600, errors='replace', cwd=r"C:\Users\Admin")
    if result.returncode == 0:
        size = os.path.getsize(output) / (1024 * 1024)
        print(f"  ✅ Pilot 合成成功: {output} ({size:.1f} MB)")
        return True
    else:
        print(f"  ❌ 合成失败")
        print(result.stderr[-500:] if result.stderr else "无错误信息")
        return False


def step6_verify():
    """验证最终输出"""
    print("\n" + "=" * 60)
    print("验证最终输出")
    print("=" * 60)

    output = os.path.join(EXPORT_DIR, "pilot_45s_v1.mp4")
    if not os.path.exists(output):
        print(f"  ❌ 文件不存在: {output}")
        return False

    size = os.path.getsize(output) / (1024 * 1024)
    print(f"  📁 {output}")
    print(f"  大小: {size:.1f} MB")

    # Get duration
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration,size",
        "-of", "default=noprint_wrappers=1:nokey=1",
        output,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
    parts = result.stdout.strip().split("\n")
    if len(parts) >= 2:
        print(f"  时长: {float(parts[1]):.1f}s")

    # Get video info
    cmd = [
        "ffprobe", "-v", "error",
        "-select_streams", "v:0",
        "-show_entries", "stream=width,height,codec_name",
        "-of", "default=noprint_wrappers=1:nokey=1",
        output,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
    parts = result.stdout.strip().split("\n")
    if len(parts) >= 3:
        print(f"  编码: {parts[2]}")
        print(f"  分辨率: {parts[0]}x{parts[1]}")

    # Check audio streams
    cmd = [
        "ffprobe", "-v", "error",
        "-select_streams", "a:0",
        "-show_entries", "stream=codec_name",
        "-of", "default=noprint_wrappers=1:nokey=1",
        output,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
    audio_codec = result.stdout.strip()
    print(f"  音频编码: {audio_codec}")

    return True


def main():
    ensure_dirs()

    print(f"Phase 4A — 45s Pilot 合成")
    print(f"总时长: {TOTAL_DURATION:.0f}s | 5 scenes")
    print()

    # Step 1: Generate TTS
    tts_ok = step1_generate_tts()
    if not tts_ok:
        print("[错误] TTS 生成失败")
        return 1

    # Step 2: Get TTS durations
    print("\nTTS 时长:")
    tts_durations = step2_get_tts_durations()

    # Step 3: Generate subtitles
    srt_path = step3_generate_subtitles(tts_durations)

    # Step 4: Generate BGM + SFX
    bgm_path, sfx_path = step4_generate_sfx()

    # Step 5: Compose final video
    ok = step5_compose_final(bgm_path, sfx_path)

    # Step 6: Verify
    verify_ok = step6_verify()

    if ok and verify_ok:
        print("\n" + "=" * 60)
        print("🎉 Phase 4A — 45s Pilot 合成完成!")
        print("=" * 60)
        return 0
    else:
        print("\n❌ Phase 4A 合成失败")
        return 1


if __name__ == "__main__":
    sys.exit(main())